import numpy as np
from scipy.optimize import least_squares
from scipy.integrate import solve_ivp
import matplotlib.pyplot as plt

# 系统参数（以Test 1为例）
g_mars = np.array([0, 0, -3.7114])  # 火星重力加速度 (m/s²)
T_min = 4971.81  # 最小推力 (N)
T_max = 13258.18  # 最大推力 (N)
alpha = 5.0863e-4  # 燃料消耗系数 (s/m)
m0 = 1905.0        # 初始质量 (kg)
r0 = np.array([-900, 10, 1500])   # 初始位置 (m)
v0 = np.array([30, -10, -70])     # 初始速度 (m/s)
rf = np.array([0, 0, 0])          # 目标位置 (m)
vf = np.array([0, 0, 0])          # 目标速度 (m/s)

# 时间参数
t0 = 0.0
tf_guess = 80.0   # 初始猜测总时间
t1_guess = 30.0   # 初始猜测切换时间

# Chebyshev多项式参数
N = 10            # 多项式阶数
n_segments = 2    # 时间分段（min-max切换）

# 定义Chebyshev基函数
def chebyshev_basis(z, n):
    return np.polynomial.chebyshev.chebvander(z, n-1)

# 映射时间到[-1,1]区间
def time_mapping(t, t_start, t_end):
    return 2*(t - t_start)/(t_end - t_start) - 1

# 约束表达式构建（单分段）
def constrained_expression(t_seg, t_start, t_end, xi, r_start, v_start, r_end, v_end):
    z = time_mapping(t_seg, t_start, t_end)
    H = chebyshev_basis(z, N)
    
    # 计算Omega函数（简化为多项式形式）
    dt = t_end - t_start
    tau = (t_seg - t_start)/dt
    Omega1 = (1 - 3*tau**2 + 2*tau**3)
    Omega2 = 3*tau**2 - 2*tau**3
    Omega3 = tau*(1 - 2*tau + tau**2)*dt
    Omega4 = tau**2*(tau - 1)*dt
    
    # 构建约束表达式
    g = H @ xi
    dg = (H @ xi) * (2/dt)  # 简化的导数（需根据实际多项式调整）
    
    r = g + Omega1*(r_start - g[0]) + Omega2*(r_end - g[-1]) + \
        Omega3*(v_start - dg[0]) + Omega4*(v_end - dg[-1])
    return r

# 定义残差函数
def residuals(params):
    # 解析参数：切换时间t1，总时间tf，各段Chebyshev系数
    t1, tf = params[:2]
    xi_segments = params[2:].reshape((n_segments, 3, N))
    
    # 时间分段
    t_segments = [np.linspace(t0, t1, 50), np.linspace(t1, tf, 50)]
    
    residuals = []
    
    # 遍历每个时间段
    for seg in range(n_segments):
        t = t_segments[seg]
        xi = xi_segments[seg]
        
        # 获取初始和终止条件
        if seg == 0:
            r_start, v_start = r0, v0
            r_end_guess = r0 + (rf - r0)*(t1/tf_guess)
            v_end_guess = v0 + (vf - v0)*(t1/tf_guess)
        else:
            r_start, v_start = r_end_guess, v_end_guess
            r_end_guess, v_end_guess = rf, vf
        
        # 计算约束表达式
        r = constrained_expression(t, t_segments[seg][0], t_segments[seg][-1], xi, 
                                  r_start, v_start, r_end_guess, v_end_guess)
        
        # 计算动力学残差（简化为位置二阶导）
        # 实际需根据论文中的状态方程和协态方程计算
        acc = np.gradient(np.gradient(r, t), t)
        T = T_min if seg == 0 else T_max  # 假设第一段最小推力
        acc_desired = g_mars + T/m0 * (-1)  # 假设推力方向
        
        residuals.extend(acc - acc_desired)
    
    return np.concatenate(residuals)

# 初始参数猜测
initial_params = np.concatenate([
    [t1_guess, tf_guess],          # 切换时间和总时间
    np.random.randn(n_segments * 3 * N)  # Chebyshev系数
])

# 使用最小二乘优化
result = least_squares(residuals, initial_params, verbose=2, max_nfev=50)

# 提取优化结果
t1_opt, tf_opt = result.x[:2]
xi_opt = result.x[2:].reshape((n_segments, 3, N))

# 生成最终轨迹
t_span = np.linspace(t0, tf_opt, 100)
r_traj = []
for seg in range(n_segments):
    t_seg = t_span[(t_span >= (t0 if seg==0 else t1_opt)) & (t_span <= (t1_opt if seg==0 else tf_opt))]
    r_seg = constrained_expression(t_seg, t_seg[0], t_seg[-1], xi_opt[seg], 
                                  r0 if seg==0 else r_mid, 
                                  v0 if seg==0 else v_mid,
                                  rf if seg==1 else r_mid,
                                  vf if seg==1 else v_mid)
    r_traj.append(r_seg)

# 可视化结果
plt.figure(figsize=(12, 6))
plt.subplot(121, projection='3d')
for seg in r_traj:
    plt.plot(seg[:,0], seg[:,1], seg[:,2], label=f'Segment {seg+1}')
plt.scatter(r0[0], r0[1], r0[2], c='r', label='Start')
plt.scatter(rf[0], rf[1], rf[2], c='g', label='Target')
plt.xlabel('X (m)'); plt.ylabel('Y (m)'); plt.title('3D Trajectory')

plt.subplot(122)
plt.plot(t_span, np.linalg.norm(r_traj, axis=1))
plt.xlabel('Time (s)'); plt.ylabel('Altitude (m)'); plt.title('Altitude vs Time')
plt.tight_layout()
plt.show()